package edu.mit.compilers.assembly.x64;

import edu.mit.compilers.assembly.x64.basic.*;
import edu.mit.compilers.assembly.x64.codes.*;
import edu.mit.compilers.ir.type.IRArrayType;
import edu.mit.compilers.ir.type.IRBasicType;
import edu.mit.compilers.lowercode.*;
import edu.mit.compilers.lowercode.code.*;
import edu.mit.compilers.utils.OperatorUtils;

import java.util.*;

// visit the lower code (three address code list) and convert it to x64-code
public class X64MethodConvertor implements ThreadAddressCodeVistor<Void> {
    private class CurrentState {
        public MethodInfo methodInfo = null;
        public X64EntryAllocator allocator = null;
        public X64EnterCode enterCode = null;

        public ArrayList<String> pushArgsInRegister = new ArrayList<>();
        public Stack<String> pushArgsInStack = new Stack<>();

        public ArrayList<String> popArgsInRegister = new ArrayList<>();
        public ArrayList<String> popArgsInStack = new ArrayList<>();

        public void resetCurrentInfo() {
            methodInfo = null;
            // allocator = null;
            enterCode = null;
            pushArgsInRegister.clear();
            pushArgsInStack.clear();
            popArgsInRegister.clear();
            popArgsInStack.clear();
        }

        public void resetPushArgsInfo() {
            pushArgsInRegister.clear();
            pushArgsInStack.clear();
        }
    }

    public static final long UNKNOWN_SIZE = -1;

    private ThreeAddressCodesInfo info;
    private X64CodeList x64CodeList;
    private final Map<String, FieldInfo>  globalFields;

    // current state
    private CurrentState ctx;

    public X64MethodConvertor(ThreeAddressCodesInfo info, X64CodeList x64CodeList) {
        this.info = info;
        this.x64CodeList = x64CodeList;
        this.globalFields = info.symbolTable.getGlobalFieldsInfo();
    }

    @Override
    public Void visit(AssignSingleOperandCode code) {
        // consider increment op like ++, --
        if (code.isIncrementOp()) {
            String incrementOp = code.getIncrementOp();
            X64Entry dest = allocateForOperand(code.getLeft(), x64CodeList);
            if (incrementOp.equals("++") || incrementOp.equals("--")) {
                X64IncrementCode.Type type = incrementOp.equals("++")
                        ? X64IncrementCode.Type.INC
                        : X64IncrementCode.Type.DEC;
                x64CodeList.append(new X64IncrementCode(dest.getLocation(), type));
                releaseUselessRegister(dest.getLocation());
                return null;
            }
        }

        // consider unary op like "-", "!" and increment op like "+=", "-="
        X64Entry right = allocateForOperand(code.getRight(), x64CodeList);
        X64Entry src = right;
        X64Entry dest = allocateForOperand(code.getLeft(), x64CodeList);

        if (code.isUnaryOp() || code.isIncrementOp()) {
            if (code.isUnaryOp()) {
                // instruction @neg just accept register
                src = allocateRegisterForEntry(right, false, x64CodeList);
                String unaryOp = code.getUnaryOp();
                if (unaryOp.equals("-")) {
                    x64CodeList.append(new X64UnaryOpCode(src.getLocation(), X64UnaryOpCode.Type.NEG));
                } else {
                    assert unaryOp.equals("!") : "unsupported unaryOp " + unaryOp;
                    X64Entry entry = ctx.allocator.allocateImmediate("$1").getEntry();
                    x64CodeList.append(constructBinaryOpCode(entry.getLocation(), src.getLocation(), X64BinaryOpCode.Type.XOR));
                }
            } else {
                src = allocateRegisterForEntry(dest, false, x64CodeList);
                String incrementOp = code.getIncrementOp();
                X64BinaryOpCode.Type type = incrementOp.equals("+=")
                        ? X64BinaryOpCode.Type.ADD
                        : X64BinaryOpCode.Type.SUB;
                x64CodeList.append(constructBinaryOpCode(right.getLocation(), src.getLocation(), type));
            }
        }

        if (src.getLocation().isMemoryOperand() && dest.getLocation().isMemoryOperand()) {
            src = allocateRegisterForEntry(src, false, x64CodeList);
        }

        x64CodeList.append(constructMoveCode(src.getLocation(), dest.getLocation()));
        return null;
    }

    @Override
    public Void visit(AssignTwoOperandCode code) {
        X64Entry dest = allocateForOperand(code.getRight1(), x64CodeList);
        X64Entry src = allocateForOperand(code.getRight2(), x64CodeList);
        String binaryOp = code.getBinaryOp();
        if (OperatorUtils.isArith(code.getBinaryOp()) || OperatorUtils.isCond(code.getBinaryOp())) {
            // div need special instruction @idivq
            if (binaryOp.equals("/") || binaryOp.equals("%")) {
                ArrayList<X64EntryAllocator.AllocateResult> res = ctx.allocator.allocateRegistersForDiv();
                X64Entry quotient = res.get(0).getEntry();  // also numerator
                X64Entry numerator = quotient;
                X64Entry remainder = res.get(1).getEntry();
                x64CodeList.append(constructMoveCode(dest.getLocation(), numerator.getLocation()));

                // maybe float exception if remainder (%rdx) is random, so need instruction @cqto;
                x64CodeList.append(new X64FlagCode("cqto", true));
                x64CodeList.append(constructSpecialOpCode(src.getLocation(), X64SpecialOpCode.Type.IDIV));
                dest = binaryOp.equals("/") ? quotient : remainder;
            } else {
                X64BinaryOpCode.Type type = null;
                switch (binaryOp) {
                    case "+":
                        type = X64BinaryOpCode.Type.ADD;
                        break;
                    case "-":
                        type = X64BinaryOpCode.Type.SUB;
                        break;
                    case "*":
                        type = X64BinaryOpCode.Type.IMUL;
                        break;
                    case "&&":
                        type = X64BinaryOpCode.Type.AND;
                        break;
                    case "||":
                        type = X64BinaryOpCode.Type.OR;
                        break;
                    default:
                        assert false : "not support binary " + code.getBinaryOp();
                }
                dest = allocateRegisterForEntry(dest, true, x64CodeList);
                x64CodeList.append(constructBinaryOpCode(src.getLocation(), dest.getLocation(), type));
            }
        } else {
            assert OperatorUtils.isRel(binaryOp) || OperatorUtils.isEq(binaryOp) : binaryOp;
            dest = allocateRegisterForEntry(dest, true, x64CodeList);

            X64SetCode.Type type = null;
            switch (binaryOp) {
                case "==":
                    type = X64SetCode.Type.EQ;
                    break;
                case "!=":
                    type = X64SetCode.Type.NE;
                    break;
                case ">=":
                    type = X64SetCode.Type.GE;
                    break;
                case ">":
                    type = X64SetCode.Type.GT;
                    break;
                case "<=":
                    type = X64SetCode.Type.LE;
                    break;
                case "<":
                    type = X64SetCode.Type.LT;
                    break;
                default:
                    assert false : "not support op " + binaryOp;
            }
            x64CodeList.append(constructCompareAndSetCodes(src.getLocation(), dest.getLocation(), type));
        }
        X64Entry left = allocateForOperand(code.getLeft(), x64CodeList);
        x64CodeList.append(constructMoveCode(dest.getLocation(), left.getLocation()));
        return null;
    }

    @Override
    public Void visit(EmptyCode code) {
        assert false : "Empty code will not exist in this step";
        return null;
    }

    @Override
    public Void visit(GotoCode code) {
        X64JumpCode.Type type = X64JumpCode.Type.NONE;
        if (code.isConditional()) {
            String conditionVar = code.getCondition();
            // entry vs zero, true means greater (ne)
            X64Entry entry = ctx.allocator.allocate(conditionVar).getEntry();
            X64Entry zero = ctx.allocator.allocateImmediate("$0").getEntry();
            x64CodeList.append(new X64CompareCode(zero.getLocation(), entry.getLocation()));
            releaseUselessRegister(entry.getLocation());
            type = code.isIfTrue() ? X64JumpCode.Type.NE : X64JumpCode.Type.EQ;
        }
        X64Entry where = ctx.allocator.allocateImmediate("." + code.getGotoLabel()).getEntry();
        x64CodeList.append(new X64JumpCode(where.getLocation(), type));
        return null;
    }

    @Override
    public Void visit(MethodBeginCode code) {
        String methodName = ctx.methodInfo.getMethodName();
        assert methodName == code.getMethodName() :
                "current method " + methodName + " is not matched with code's " + code.getMethodName();
        ctx.enterCode = new X64EnterCode(UNKNOWN_SIZE);
        x64CodeList.append(ctx.enterCode);
        return null;
    }

    // will push args into registers and stack
    // previous 6 args : %rdi %rsi %rdx %rcx %r8 %r9, others in stack in reversed order
    @Override
    public Void visit(MethodCallCode code) {
        List<X64Register> registerUsedForCaller = new ArrayList<>();
        for (int i = 0 ; i < ctx.pushArgsInRegister.size(); ++i) {
            String argName = ctx.pushArgsInRegister.get(i);
            X64Entry src = allocateForVariable(argName, x64CodeList);
            X64Entry dest = ctx.allocator.allocateArgRegister(i, 8).getEntry();
            assert dest.getLocation().isRegisterOperand();
            registerUsedForCaller.addAll(dest.getLocation().getUsedRegisters());
            x64CodeList.append(constructMoveCode(src.getLocation(), dest.getLocation()));
        }

        // use appends to save push codes and calc pop codes after call
        X64CodeList appends = new X64CodeList();
        long stackNeedMoreSize = 0;

        // push args in reversed order
        while (!ctx.pushArgsInStack.empty()) {
            String argName = ctx.pushArgsInStack.pop();
            X64Entry entry = allocateRegisterForVariable(argName, appends);
            appends.append(constructPushCode(entry.getLocation()));
            stackNeedMoreSize += entry.size;
        }

        // align stack pointer (%rsp) to 16-aligned
        if (stackNeedMoreSize % 16 != 0) {
            X64Entry entry = ctx.allocator.allocateImmediate("$0").getEntry();
            X64PushCode alignedCode = constructPushCode(entry.getLocation());

            // appends = aligned code || appends
            appends = new X64CodeList().append(alignedCode).append(appends);
            stackNeedMoreSize += 8;
        }
        x64CodeList.append(appends);

        // call method, and store return value if necessary
        X64Entry methodEntry = ctx.allocator.allocateImmediate(code.getMethodName()).getEntry();
        x64CodeList.append(new X64CallCode(methodEntry.getLocation()));
        if (code.isAssign()) {
            X64Entry srcEntry = ctx.allocator.allocateRetRegister(8).getEntry();
            X64Entry destEntry = allocateForVariable(code.getVariable(), x64CodeList);
            x64CodeList.append(constructMoveCode(srcEntry.getLocation(), destEntry.getLocation()));
        }

        for (X64Register register : registerUsedForCaller) {
            ctx.allocator.releaseTempRegister(register);
        }
        // restore stack pointer at before
        if (stackNeedMoreSize > 0) {
            X64Entry addEntry = ctx.allocator.allocateImmediate('$' + String.valueOf(stackNeedMoreSize)).getEntry();
            X64Entry stackPointEntry = ctx.allocator.allocateRunStackPointer().getEntry();
            x64CodeList.append(constructBinaryOpCode(addEntry.getLocation(), stackPointEntry.getLocation(), X64BinaryOpCode.Type.ADD));
        }
        ctx.resetPushArgsInfo();
        return null;
    }

    @Override
    public Void visit(MethodLeaveCode code) {
        if (ctx.methodInfo.getReturnType() == IRBasicType.VoidType) {
            X64Entry src = ctx.allocator.allocateImmediate("$0").getEntry();
            X64Entry dest = ctx.allocator.allocateRetRegister(8).getEntry();
            x64CodeList.append(constructMoveCode(src.getLocation(), dest.getLocation()));
        }
        x64CodeList.append(X64FlagCode.LEAVE_CODE);
        x64CodeList.append(X64FlagCode.RET_CODE);
        assert ctx.enterCode != null;
        ctx.enterCode.setStackFrameSize(ctx.allocator.getStackFrameSize());
        ctx.resetCurrentInfo();
        return null;
    }

    @Override
    public Void visit(MethodSetResultCode code) {
        String variable = code.getVariable();
        X64Operand src = ctx.allocator.allocate(variable).getEntry().getLocation();
        X64Operand dest = ctx.allocator.allocateRetRegister(8).getEntry().getLocation();
        x64CodeList.append(constructMoveCode(src, dest));
        return null;
    }

    @Override
    public Void visit(PopParamCode code) {
        if (ctx.popArgsInRegister.size() < 6) {
            int i = ctx.popArgsInRegister.size();
            ctx.popArgsInRegister.add(code.getVariable());
            X64Entry src = ctx.allocator.allocateArgRegister(i, 8).getEntry();
            X64Entry dest = allocateForVariable(code.getVariable(), x64CodeList);
            assert src.getLocation().isRegisterOperand();
            x64CodeList.append(constructMoveCode(src.getLocation(), dest.getLocation()));
        } else {
            int i = ctx.popArgsInStack.size();
            ctx.popArgsInStack.add(code.getVariable());
            X64Entry src = allocateRegisterForEntry(ctx.allocator.allocatePopArgInStack(i).getEntry(), false, x64CodeList);
            X64Entry dest = ctx.allocator.allocate(code.getVariable()).getEntry();
            x64CodeList.append(constructMoveCode(src.getLocation(), dest.getLocation()));
        }
        return null;
    }

    @Override
    public Void visit(PushParamCode code) {
        if (ctx.pushArgsInRegister.size() < 6) {
            ctx.pushArgsInRegister.add(code.getVariable());
        } else {
            ctx.pushArgsInStack.add(code.getVariable());
        }
        return null;
    }

    @Override
    public Void visit(ThrowRuntimeErrorCode code) {
        X64Entry method = ctx.allocator.allocateGlobal(X64Util.builtInMethodName(ctx.methodInfo.getMethodName()), 8).getEntry();
        X64Entry rsi = ctx.allocator.allocateArgRegister(1, 8).getEntry();
        x64CodeList.append(new X64LeaCode(method.getLocation(), rsi.getLocation()));

        // if out of bound, goto exception handler
        X64Entry label = ctx.allocator.allocateImmediate(X64RuntimeChecker.errorHandleLabel(code.getEc())).getEntry();
        x64CodeList.append(new X64JumpCode(label.getLocation(), X64JumpCode.Type.NONE));
        releaseUselessRegister(rsi.getLocation());
        return null;
    }

    public void convertMethod(String methodName) {
        ThreeAddressCodeList codes = info.method2Codes.get(methodName);
        assert codes != null : "method " + methodName + " not found in method2Codes";
        assert !codes.isNull() : "method " + methodName + " has no codes";

        ctx = new CurrentState();
        ctx.methodInfo = info.symbolTable.getMethodInfo(methodName);
        ctx.allocator = new X64EntryAllocator(methodName, info.symbolTable, codes);

        for (ThreeAddressCode code : codes) {
            if (code.hasLabel()) {
                x64CodeList.append(new X64LabelCode("." + code.getLabel()));
            }
            x64CodeList.append(new X64DebugInfoCode(code));
            code.accept(this);
        }
    }

    public void finish() {
        X64RuntimeChecker.fillErrorHandlers(x64CodeList, ctx.allocator);
    }

    // must be simple variable or const string variable
    private X64Entry allocateForVariable(String name, X64CodeList appender) {
        assert name != null;
        SymbolTable symbolTable = info.symbolTable;
        if (symbolTable.isGlobalFieldName(name) && symbolTable.getGlobalFieldInfo(name).getType().isBasicType()) {
            return ctx.allocator.allocateGlobal(name, 8).getEntry();
        }

        FieldInfo fieldInfo = getFieldInfo(name);
        if (fieldInfo.getType().isArrayType() || fieldInfo.getVariable().hasConstValue()) {
            assert appender != null;
            X64Entry src = fieldInfo.getVariable().isGlobal()
                    ? ctx.allocator.allocateGlobal(name, 8).getEntry()
                    : ctx.allocator.allocate(name).getEntry();
            X64Entry allocated = ctx.allocator.allocateTempRegister(8).getEntry();
            appender.append(constructLeaCode(src.getLocation(), allocated.getLocation()));
            return allocated;
        }

        return ctx.allocator.allocate(name).getEntry();
    }

    private FieldInfo getFieldInfo(String name) {
        if (ctx.methodInfo != null && ctx.methodInfo.getFieldsInfo().containsKey(name)) {
            return ctx.methodInfo.getFieldsInfo().get(name);
        }
        if (info.symbolTable.isGlobalFieldName(name)) {
            return info.symbolTable.getGlobalFieldInfo(name);
        }
        return info.symbolTable.getConstStringInfo().getOrDefault(name, null);
    }

    private X64Entry allocateForValue(Value value) {
        return ctx.allocator.allocateImmediate('$' + value.getLiteralString()).getEntry();
    }

    private X64Entry allocateForLocationValue(LocationVariable locationVariable, X64CodeList appender) {
        assert locationVariable.getLocation() instanceof Variable : locationVariable.getLiteralString() + " some wrong";
        String variable = locationVariable.getVariable().getName();
        String index = ((Variable) locationVariable.getLocation()).getName();

        X64Entry varEntry = allocateRegisterForVariable(variable, appender);
        X64Entry indexEntry = allocateRegisterForVariable(index, appender);

        String baseRegister = varEntry.getLocation().getUsedRegisters().get(0).name;
        String offsetRegister = indexEntry.getLocation().getUsedRegisters().get(0).name;

        FieldInfo fieldInfo = getFieldInfo(variable);
        int scaled = (int) ((IRArrayType) fieldInfo.getType()).getBasicType().getSize();
        assert scaled == 8;
        appender.append(constructArrayCheckCodes(indexEntry, fieldInfo.getSize() / scaled));
        return ctx.allocator.allocateScaled(baseRegister, offsetRegister, scaled).getEntry();
    }

    private X64Entry allocateRegisterForEntry(X64Entry entry, boolean forceTemp, X64CodeList appender) {
        X64Entry ret = entry;
        X64Operand location = entry.getLocation();
        if (!location.isRegisterOperand() || (forceTemp && !location.getUsedRegisters().get(0).temporary)) {
            X64Entry tmpRegister = ctx.allocator.allocateTempRegister(8).getEntry();
            appender.append(constructMoveCode(location, tmpRegister.getLocation()));
            ret = tmpRegister;
        }
        return ret;
    }

    private X64Entry allocateRegisterForVariable(String variable, X64CodeList appender) {
        X64Entry ret = allocateForVariable(variable, appender);
        return allocateRegisterForEntry(ret, false, appender);
    }

    // operand not constString
    private X64Entry allocateForOperand(Operand operand, X64CodeList appender) {
        if (operand instanceof Value) {
            return allocateForValue((Value) operand);
        }

        if (operand instanceof LocationVariable && ((LocationVariable) operand).getLocation() != null) {
            return allocateForLocationValue((LocationVariable) operand, appender);
        }
        return allocateForVariable(operand.getLiteralString(), appender);
    }

    private void releaseUselessRegister(X64Operand src, X64Operand dest) {
        // we assume that temp register's ref count will less or equal than 1 (only 1 owner at most)
        for (X64Register register : src.getUsedRegisters()) {
            if (!dest.isRegisterOperand() || (!dest.getUsedRegisters().contains(register) && register.temporary)) {
                ctx.allocator.releaseTempRegister(register);
            }
        }
        if (dest.isMemoryOperand()) {
            for (X64Register register : dest.getUsedRegisters()) {
                ctx.allocator.releaseTempRegister(register);
            }
        }
    }

    private void releaseUselessRegister(X64Operand src) {
        releaseUselessRegister(src, new X64ImmOperand("1"));
    }

    private X64MoveCode constructMoveCode(X64Operand src, X64Operand dest) {
        releaseUselessRegister(src, dest);
        return new X64MoveCode(src, dest);
    }

    private X64PushCode constructPushCode(X64Operand src) {
        releaseUselessRegister(src);
        return new X64PushCode(src);
    }

    private X64LeaCode constructLeaCode(X64Operand src, X64Operand dest) {
        releaseUselessRegister(src, dest);
        return new X64LeaCode(src, dest);
    }

    private X64BinaryOpCode constructBinaryOpCode(X64Operand src, X64Operand dest, X64BinaryOpCode.Type type) {
        releaseUselessRegister(src, dest);
        return new X64BinaryOpCode(src, dest, type);
    }

    private X64SpecialOpCode constructSpecialOpCode(X64Operand src, X64SpecialOpCode.Type type) {
        releaseUselessRegister(src);
        return new X64SpecialOpCode(src, type);
    }

    private X64CodeList constructCompareAndSetCodes(X64Operand src, X64Operand dest, X64SetCode.Type type) {
        X64CodeList ret = new X64CodeList().append(new X64CompareCode(src, dest));
        X64Register register = X64RegisterFamily.getRegister(dest.getUsedRegisters().get(0).name, 1);
        X64Operand registerOperand = new X64RegisterOperand(register);
        ret.append(new X64SetCode(registerOperand, type));
        ret.append(new X64MovezCode(registerOperand, dest));
        releaseUselessRegister(src);
        return ret;
    }

    private X64CodeList constructArrayCheckCodes(X64Entry entry, long bound) {
        X64CodeList ret = new X64CodeList();
        X64Entry zero = ctx.allocator.allocateImmediate("$0").getEntry();
        X64Entry upperBound = ctx.allocator.allocateImmediate("$" + bound).getEntry();

        // check 0 <= entry < bound
        X64Entry tmp1 = ctx.allocator.allocateTempRegister(8).getEntry();
        ret.append(new X64MoveCode(entry.getLocation(), tmp1.getLocation()));
        ret.append(constructCompareAndSetCodes(zero.getLocation(), tmp1.getLocation(), X64SetCode.Type.LT));

        X64Entry tmp2 = ctx.allocator.allocateTempRegister(8).getEntry();
        ret.append(new X64MoveCode(entry.getLocation(), tmp2.getLocation()));
        ret.append(constructCompareAndSetCodes(upperBound.getLocation(), tmp2.getLocation(), X64SetCode.Type.GE));
        ret.append(constructBinaryOpCode(tmp1.getLocation(), tmp2.getLocation(), X64BinaryOpCode.Type.OR));

        ret.append(new X64CompareCode(zero.getLocation(), tmp2.getLocation()));
        releaseUselessRegister(tmp2.getLocation());
        X64Entry method = ctx.allocator.allocateGlobal(X64Util.builtInMethodName(ctx.methodInfo.getMethodName()), 8).getEntry();
        X64Entry rsi = ctx.allocator.allocateArgRegister(1, 8).getEntry();
        ret.append(new X64PushCode(zero.getLocation()));
        ret.append(new X64PushCode(rsi.getLocation()));
        ret.append(new X64LeaCode(method.getLocation(), rsi.getLocation()));

        // if out of bound, goto exception handler
        X64Entry label = ctx.allocator.allocateImmediate(X64RuntimeChecker.errorHandleLabel(1)).getEntry();
        ret.append(new X64JumpCode(label.getLocation(), X64JumpCode.Type.NE));
        ret.append(new X64PopCode(rsi.getLocation()));
        ret.append(new X64BinaryOpCode(
                ctx.allocator.allocateImmediate("$8").getEntry().getLocation(),
                ctx.allocator.allocateRunStackPointer().getEntry().getLocation(),
                X64BinaryOpCode.Type.ADD));
        releaseUselessRegister(tmp1.getLocation());
        return ret;
    }
}
